// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <pollux/common/caching/file_ids.h>
#include <pollux/common/caching/scan_tracker.h>
#include <pollux/common/caching/ssd_cache.h>
#include <pollux/common/io/io_statistics.h>
#include <pollux/dwio/common/input_stream.h>
#include <pollux/dwio/common/seekable_input_stream.h>

namespace kumo::pollux::dwio::common {
    class CachedBufferedInput;

    class CacheInputStream : public SeekableInputStream {
    public:
        CacheInputStream(
            CachedBufferedInput *cache,
            IoStatistics *ioStats,
            const pollux::common::Region &region,
            std::shared_ptr<ReadFileInputStream> input,
            uint64_t fileNum,
            bool noCacheRetention,
            std::shared_ptr<cache::ScanTracker> tracker,
            cache::TrackingId trackingId,
            uint64_t groupId,
            int32_t loadQuantum);

        ~CacheInputStream() override;

        CacheInputStream &operator=(const CacheInputStream &) = delete;

        CacheInputStream(const CacheInputStream &) = delete;

        CacheInputStream &operator=(CacheInputStream &&) = delete;

        CacheInputStream(CacheInputStream &&) = delete;

        bool Next(const void **data, int *size) override;

        void BackUp(int count) override;

        bool SkipInt64(int64_t count) override;

        google::protobuf::int64 ByteCount() const override;

        void seekToPosition(PositionProvider &position) override;

        std::string getName() const override;

        size_t positionSize() const override;

        /// Returns a copy of 'this', ranging over the same bytes. The clone is
        /// initially positioned at the position of 'this' and can be moved
        /// independently within 'region_'.  This is used for first caching a range of
        /// bytes from a file and then giving out delimited subranges of this to
        /// callers. skip() and setRemainingBytes() set the bounds of the window
        /// exposed by the clone. In specific, reading protocol buffers requires a
        /// stream that begins and ends at the exact start and end of the
        /// serialization. Reading these from cache requires an exactly delimited
        /// stream.
        std::unique_ptr<CacheInputStream> clone() {
            auto copy = std::make_unique<CacheInputStream>(
                bufferedInput_,
                ioStats_,
                region_,
                input_,
                fileNum_,
                noCacheRetention_,
                tracker_,
                trackingId_,
                groupId_,
                loadQuantum_);
            copy->position_ = position_;
            return copy;
        }

        /// Sets the stream to range over a window that starts at the current position
        /// and is 'remainingBytes' bytes in size. 'remainingBytes' must be <=
        /// 'region_.length - position_'. The stream cannot be used for reading
        /// outside of the window. Use together with clone() and skip().
        void setRemainingBytes(uint64_t remainingBytes);

        /// Causes the next load quantum to be scheduled for read-ahead when
        /// 'percent' of the current load quantum has been returned by
        /// Next(). If Next() returns the whole read quantum, them the first
        /// Next triggers the read ahead of te next quantum right away. a
        /// value of over 100 causes no prefetches to be made. If there is
        /// no memory to cover the load quantum to prefetch the prefetch
        /// fails silently.
        void setPrefetchPct(int32_t pct) {
            prefetchPct_ = pct;
        }

        bool testingNoCacheRetention() const {
            return noCacheRetention_;
        }

    private:
        // Ensures that the current position is covered by 'pin_'.
        void loadPosition();

        // Returns the next quantized region to load with given loaded position.
        pollux::common::Region nextQuantizedLoadRegion(
            uint64_t prevLoadedPosition) const;

        // Synchronously sets 'pin_' to cover 'region'.
        void loadSync(const pollux::common::Region &region);

        // Returns true if there is an SSD cache and 'entry' is present there and
        // successfully loaded.
        bool loadFromSsd(
            const pollux::common::Region &region,
            cache::AsyncDataCacheEntry &entry);

        // Invoked to clear the cache pin of the accessed cache entry and mark it as
        // immediate evictable if 'noCacheRetention_' flag is set.
        void clearCachePin();

        void makeCacheEvictable();

        // Return SSD cache file path if exists; return empty string if no SSD cache
        // file.
        std::string ssdFileName() const;

        CachedBufferedInput *const bufferedInput_;
        cache::AsyncDataCache *const cache_;
        // True if a pin should be set to the lowest retention score after
        // unpinning. This applies to sequential reads where second access
        // to the page is not expected.
        const bool noCacheRetention_;
        // The region of 'input' 'this' ranges over.
        const pollux::common::Region region_;
        const uint64_t fileNum_;
        const std::shared_ptr<cache::ScanTracker> tracker_;
        const cache::TrackingId trackingId_;
        const uint64_t groupId_;

        // Maximum number of bytes read from 'input' at a time. This gives the maximum
        // pin_.entry()->size().
        const int32_t loadQuantum_;

        IoStatistics *const ioStats_;
        const std::shared_ptr<ReadFileInputStream> input_;

        // Handle of cache entry.
        cache::CachePin pin_;

        // Offset of current run from start of 'entry_->data()'
        uint64_t offsetOfRun_;

        // Pointer  to start of  current run in 'entry->data()' or
        // 'entry->tinyData()'.
        uint8_t *run_{nullptr};
        // Position of stream relative to 'run_'.
        int offsetInRun_{0};
        // Index of run in 'entry_->data()'
        int runIndex_ = -1;
        // Number of valid bytes above 'run_'.
        uint32_t runSize_ = 0;
        // Position relative to 'region_.offset'.
        uint64_t position_ = 0;

        // A restricted view over 'region'. offset is relative to 'region_'. A cloned
        // CacheInputStream can cover a sub-range of the range of the original.
        std::optional<pollux::common::Region> window_;

        // Percentage of 'loadQuantum_' at which the next load quantum gets scheduled.
        // Over 100 means no prefetch.
        int32_t prefetchPct_{200};

        // True if prefetch the next 'loadQuantum_' has been started. Cleared when
        // moving to the next load quantum.
        bool prefetchStarted_{false};
    };
} // namespace kumo::pollux::dwio::common
